#!/usr/bin/env python3
import sys
import math
import os, time, random
import json
from multiprocessing import Pool

def back_up_origin_data(src_task_file, processed_rest_dir):
    fp = open(src_task_file, "r")
    for line in fp.readlines():
        subj_name = line.replace("\n", "")
        subj_name = subj_name.strip()

        sud_dir = os.path.join(processed_rest_dir, subj_name)
        surf_dir = os.path.join(sud_dir, 'surf')
        surf_qc_dir = os.path.join(sud_dir, 'surf_origin')

        if os.path.exists(surf_qc_dir):
            rm_cmd = 'rm -rf ' + surf_dir
            print('[CMD]: %s' % (rm_cmd))
            os.system(rm_cmd)

            mv_cmd = 'mv ' + surf_qc_dir + ' ' + surf_dir
            print('[CMD]: %s' % (mv_cmd))
            os.system(mv_cmd)

        cp_cmd = 'cp -R ' + surf_dir + ' ' + surf_qc_dir
        print('#' * 30)
        print('[CMD]: %s\n' % cp_cmd)
        os.system(cp_cmd)

def check_diff(src_task_file, processed_rest_dir):
    print("=============== Checking file =============")

    fp = open(src_task_file, "r")
    for line in fp.readlines():
        subj_name = line.replace("\n", "")
        subj_name = subj_name.strip()

        sud_dir = os.path.join(processed_rest_dir, subj_name)
        surf_dir = os.path.join(sud_dir, 'surf')
        surf_qc_dir = os.path.join(sud_dir, 'surf_origin')

        surf_list = os.listdir(surf_dir)
        surf_qc_list = os.listdir(surf_qc_dir)

        surf_sess_list = extract_sess_list(surf_dir)
        surf_qc_sess_list = extract_sess_list(surf_qc_dir)

        if surf_list != surf_qc_list:
            print('[WARNING]: Check file different in : `%s`, Sessions: %s' % (surf_dir, ','.join(surf_sess_list)))
            print('[WARNING]: Check file different in : `%s`, Sessions: %s' % (surf_qc_dir, ','.join(surf_qc_sess_list)))

    print("================== Done ==================")

def rename_file(src_task_file, processed_rest_dir):
    print("=============== Rename file =============")

    fp = open(src_task_file, "r")
    for line in fp.readlines():
        subj_name = line.replace("\n", "")
        subj_name = subj_name.strip()

        sud_dir = os.path.join(processed_rest_dir, subj_name)
        surf_dir = os.path.join(sud_dir, 'surf')

        surf_list = os.listdir(surf_dir)

        if len(surf_list) > 0:
            for surf_file_item in surf_list:
                if 'fs6_sm6_fs4' in surf_file_item:
                    tar_file = surf_file_item.replace('fs6_sm6_fs4', 'fsaverage6_sm6_fsaverage4')
                    full_tar_file = os.path.join(surf_dir, tar_file)

                    full_src_file = os.path.join(surf_dir, surf_file_item)

                    mv_cmd = 'cp ' + full_src_file + ' ' + full_tar_file
                    print('[CMD]: %s' % (mv_cmd))
                    os.system(mv_cmd)

    print("================== Done ==================")

def extract_sess_list(surf_dir):
    surf_list = os.listdir(surf_dir)

    sess_list = []
    for surf_item in surf_list:
        sess_id = surf_item.split('_')[2]
        sess_list.append(sess_id)

    sess_list = list(set(sess_list))
    sess_list.sort()
    return sess_list

def censor_proper_files(src_task_file, processed_rest_dir, remain_dict):
    subj_list = []
    exclude_list = []

    fp = open(src_task_file, "r")
    for line in fp.readlines():
        subj_name = line.replace("\n", "")

        sud_dir = os.path.join(processed_rest_dir, subj_name)
        surf_dir = os.path.join(sud_dir, 'surf')
        surf_qc_dir = os.path.join(sud_dir, 'surf_origin')

        file_list = os.listdir(surf_dir)

        remain_file_list = []
        for sess_item in remain_dict[subj_name]:
            if len(remain_dict[subj_name]) > 0:
                subj_list.append(subj_name)

                for file_item in file_list:
                    if sess_item in file_item:
                        remain_file_list.append(file_item)

        del_file_list = [file_item for file_item in file_list if not file_item in remain_file_list]
        for del_file in del_file_list:
            tar_del_file = os.path.join(surf_dir, del_file)
            if os.path.exists(tar_del_file):
                rm_cmd = 'rm ' + tar_del_file
                print('[CMD]: %s\n' % rm_cmd)
                os.system(rm_cmd)

        if len(remain_dict[subj_name]) == 0:
            rm_cmd = 'rm -rf ' + surf_dir + '/*'
            print('[CMD]: %s\n' % rm_cmd)
            os.system(rm_cmd)

            exclude_list.append(subj_name)

    subj_list = list(set(subj_list))
    exclude_list = list(set(exclude_list))

    subj_list.sort()
    exclude_list.sort()

    return subj_list, exclude_list

def save_json(remian_dict, trg_path):
    remain_json = json.dumps(remian_dict)

    fp = open(trg_path, 'w')
    fp.write(remain_json)
    fp.close()

def load_json(trg_path):
    fp = open(trg_path, 'r')
    remain_json = fp.read()
    fp.close()

    remain_dict = json.loads(remain_json)

    return remain_dict

def write_list_record(dir_list, full_file_path):
    if not os.path.exists(full_file_path):
        fp = open(full_file_path, "w")
        for sub_item in dir_list:
            fp.write(sub_item + '\n')
        fp.close()

    print("[CLEAN UP]: Writing record in `%s`" % (full_file_path))


if __name__=='__main__':
    project_dir = '/home/liang/Projects/ASD_QC'
    symbol      = "OHSU"
    sub_type    = "HC"
    list_file   = 'variability_list_' + sub_type + '_' + symbol + '.txt'

    list_dir = os.path.join(project_dir, 'Lists')
    src_task_file = os.path.join(list_dir, list_file)
    processed_rest_dir = os.path.join(project_dir, 'DataProcessed/rest')

    remain_dict = load_json('extract_{}_{}.json'.format(symbol, sub_type))
    back_up_origin_data(src_task_file, processed_rest_dir)
    check_diff(src_task_file, processed_rest_dir)

    sub_list, exclude_list = censor_proper_files(src_task_file, processed_rest_dir, remain_dict)
    check_diff(src_task_file, processed_rest_dir)

    list_file = os.path.join(list_dir, 'censored_sub_list_{}_{}.txt'.format(sub_type, symbol))
    write_list_record(sub_list, list_file)
    list_file = os.path.join(list_dir, 'censor_exclude_sub_list_{}_{}.txt'.format(sub_type, symbol))
    write_list_record(exclude_list, list_file)

    rename_file(src_task_file, processed_rest_dir)
